from typing import List

import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import xavier_init

from mmdet.models.builder import NECKS


@NECKS.register_module()
class LayerAggregation(nn.Module):
    def __init__(self, in_channels_list: List[int], out_channels: int):
        super().__init__()
        assert len(in_channels_list) >= 2
        self.lateral_convs = nn.ModuleList()
        for in_channels in in_channels_list:
            l_conv = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, bias=False),
                nn.BatchNorm2d(out_channels), nn.ReLU(inplace=False))
            self.lateral_convs.append(l_conv)
        self.conv = nn.Conv2d(out_channels, out_channels, 3, 1, 1)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def init_weights(self):
        """Initialize the weights of FPN module."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution="uniform")

    def forward(self, inputs):
        laterals = [
            lateral_conv(inputs[i])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]
        for i in range(len(inputs) - 1, 0, -1):
            laterals[i - 1] += F.interpolate(laterals[i], scale_factor=2)
        output = self.relu(self.bn(self.conv(laterals[0])))
        return output
